{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "44a00f5a-2551-42fb-b1de-272fb6e09da8",
      "metadata": {},
      "source": [
        "# Code World Model Debugger Demo\n",
        "\n",
        "In this demo, we use CWM's execution tracing capabilities to build an interactive debugger for Python programs.\n",
        "\n",
        "To get started, spin up an `fgserve` instance. For example, using four GPUs:\n",
        "```bash\n",
        "torchrun --nproc-per-node 4 \\\n",
        "    -m serve.fgserve \\\n",
        "    gen_args.host_cache_gb=80 \\\n",
        "    config=serve/configs/cwm.yaml \\\n",
        "    checkpoint_dir=<path/to/checkpoint>\n",
        "```\n",
        "\n",
        "Afterwards, update the variables in the next cell and run the notebook to get a small UI.\n",
        "Adjust the contents of the `code` variable as you like.\n",
        "\n",
        "A few notes on usage:\n",
        "- `gen_args.host_cache_gb` will speed up inference but is not required\n",
        "- \"Reset\" will reset the debugger, prompting the model to produce a new initial step\n",
        "- \"Step Out\" will prompt the model repeatedly until it predicts a return from the current function\n",
        "- You can change local variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "84f1762c-ab80-4e6d-a666-a2402a3cbff2",
      "metadata": {},
      "outputs": [],
      "source": [
        "# Point these to an fgserve instance\n",
        "FGSERVE_HOST = \"localhost\"\n",
        "FGSERVE_PORT = 5678\n",
        "\n",
        "# Point this to the CWM tokenizer file\n",
        "TOKENIZER_PATH = \"/path/to/checkpoint/tokenizer.model\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a40fcf8e-fbc6-4fdd-81e6-5693dfeb0901",
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "from pathlib import Path\n",
        "\n",
        "cwm_dir = Path(os.getcwd()).parent\n",
        "sys.path.append(str(cwm_dir))\n",
        "from cwm.text.tokenizers import build_tokenizer\n",
        "from demos.cwmdbg import CWMDebugger, CWMTraceEvent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7bd83c07-8ff3-47d5-bef9-03fae55d2de3",
      "metadata": {},
      "outputs": [],
      "source": [
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "from pygments import highlight\n",
        "from pygments.formatters import HtmlFormatter\n",
        "from pygments.lexers import PythonLexer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3caf75eb-3937-43e4-abb1-d0a306d53ccc",
      "metadata": {
        "editable": true,
        "slideshow": {
          "slide_type": ""
        },
        "tags": []
      },
      "outputs": [],
      "source": [
        "class InteractiveDebugger:\n",
        "    def __init__(self, code: str, host: str = FGSERVE_HOST, port: int = FGSERVE_PORT):\n",
        "        tokenizer = build_tokenizer(\n",
        "            name=\"cwm_instruct\",\n",
        "            path=TOKENIZER_PATH,\n",
        "        )\n",
        "        self.dbg = CWMDebugger(code, tokenizer, host, port, temperature=0.6)\n",
        "\n",
        "        self._setup_ui()\n",
        "\n",
        "    def _setup_ui(self) -> None:\n",
        "        # Setup widgets\n",
        "        self.text_area_src = widgets.HTML(\n",
        "            value=\"\",\n",
        "            disabled=True,\n",
        "            layout=widgets.Layout(width=\"50%\", height=\"400px\"),\n",
        "            style={\"font_family\": \"monospace\"},\n",
        "        )\n",
        "        self.text_area_vars = widgets.Textarea(\n",
        "            value=\"\",\n",
        "            disabled=False,\n",
        "            layout=widgets.Layout(width=\"50%\", height=\"400px\"),\n",
        "            style={\"font_family\": \"monospace\"},\n",
        "        )\n",
        "        self.next_button = widgets.Button(\n",
        "            description=\"Next\",\n",
        "            button_style=\"info\",\n",
        "        )\n",
        "        self.step_button = widgets.Button(\n",
        "            description=\"Step\",\n",
        "            button_style=\"info\",\n",
        "        )\n",
        "        self.step_out_button = widgets.Button(\n",
        "            description=\"Step Out\",\n",
        "            button_style=\"info\",\n",
        "        )\n",
        "        self.step_back_button = widgets.Button(\n",
        "            description=\"Step Back\",\n",
        "            button_style=\"info\",\n",
        "        )\n",
        "        self.reset_button = widgets.Button(\n",
        "            description=\"Reset\",\n",
        "            button_style=\"danger\",\n",
        "        )\n",
        "\n",
        "        displays = widgets.HBox([self.text_area_src, self.text_area_vars])\n",
        "        buttons = widgets.HBox(\n",
        "            [\n",
        "                self.next_button,\n",
        "                self.step_button,\n",
        "                self.step_out_button,\n",
        "                self.step_back_button,\n",
        "                self.reset_button,\n",
        "            ]\n",
        "        )\n",
        "        self.ui_container = widgets.VBox([buttons, displays])\n",
        "\n",
        "        # Setup callbacks\n",
        "        self.next_button.on_click(self._on_next_clicked)\n",
        "        self.step_button.on_click(self._on_step_clicked)\n",
        "        self.step_out_button.on_click(self._on_step_out_clicked)\n",
        "        self.step_back_button.on_click(self._on_step_back_clicked)\n",
        "        self.reset_button.on_click(self._on_reset_clicked)\n",
        "\n",
        "    # UI action decorator that handles setting variables, updating UI elements,\n",
        "    # and catching exceptions\n",
        "    def _ui_action(func):\n",
        "        import traceback\n",
        "\n",
        "        def inner(self, _):\n",
        "            try:\n",
        "                self._apply_vars()\n",
        "                func(self, _)\n",
        "                self._update_ui()\n",
        "            except Exception:\n",
        "                self.text_area_vars.value = traceback.format_exc()\n",
        "\n",
        "        return inner\n",
        "\n",
        "    # UI callbacks\n",
        "    @_ui_action\n",
        "    def _on_next_clicked(self, _):\n",
        "        self.dbg.next()\n",
        "\n",
        "    @_ui_action\n",
        "    def _on_step_clicked(self, _):\n",
        "        self.dbg.step()\n",
        "\n",
        "    @_ui_action\n",
        "    def _on_step_out_clicked(self, _):\n",
        "        self.dbg.step_out()\n",
        "\n",
        "    @_ui_action\n",
        "    def _on_step_back_clicked(self, _):\n",
        "        self.dbg.step_back()\n",
        "\n",
        "    @_ui_action\n",
        "    def _on_reset_clicked(self, _):\n",
        "        self.dbg.reset()\n",
        "\n",
        "    def _update_ui(self):\n",
        "        frame = self.dbg.current_frame\n",
        "\n",
        "        # Update source code, highlight current line\n",
        "        formatter = HtmlFormatter(\n",
        "            style=\"friendly\",\n",
        "            linenos=\"table\",\n",
        "            hl_lines=[frame.line_no] if frame.line_no else [],\n",
        "        )\n",
        "        css_styles = formatter.get_style_defs(\".highlight\")\n",
        "        html = highlight(\n",
        "            self.dbg.source,\n",
        "            PythonLexer(stripnl=False),\n",
        "            formatter,\n",
        "        )\n",
        "        full_html = f'<style>{css_styles}</style><div class=\"highlight\">{html}</div>'\n",
        "        self.text_area_src.value = full_html\n",
        "\n",
        "        # Update variables\n",
        "        if frame.event in (CWMTraceEvent.RETURN, CWMTraceEvent.EXCEPTION):\n",
        "            self.text_area_vars.value = frame.arg\n",
        "        else:\n",
        "            var_str = \"\"\n",
        "            for k, v in self.dbg.local_vars.items():\n",
        "                var_str += f\"{k}: {v}\\n\"\n",
        "            self.text_area_vars.value = var_str\n",
        "        self.text_area_vars._orig_value = self.text_area_vars.value\n",
        "\n",
        "    def _apply_vars(self):\n",
        "        if self.text_area_vars.value == self.text_area_vars._orig_value:\n",
        "            # No change\n",
        "            return\n",
        "\n",
        "        # Apply variables from input to current frame\n",
        "        lv = {}\n",
        "        for ln in self.text_area_vars.value.splitlines():\n",
        "            try:\n",
        "                k, v = ln.split(\": \", maxsplit=1)\n",
        "            except Exception:\n",
        "                return  # invalid format, ignore\n",
        "            lv[k] = v\n",
        "        self.dbg.current_frame.local_vars = lv\n",
        "        # Flush cache for subsequent prompting\n",
        "        self.dbg.current_frame._tokens = None\n",
        "\n",
        "    def display(self) -> None:\n",
        "        self.dbg.reset()\n",
        "        self._update_ui()\n",
        "        display(self.ui_container)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a52222f7-15d9-4540-81ec-6d323f84bd1e",
      "metadata": {},
      "outputs": [],
      "source": [
        "# This is the code we want to debug.\n",
        "# \"<< START_OF_TRACE\" marks the entry point of the trace and should point to a function\n",
        "code = \"\"\"\\\n",
        "def count_letters(s, letter):\n",
        "    n = 0\n",
        "    for c in s:\n",
        "        n += int(c == letter)\n",
        "    return n\n",
        "\n",
        "def format_answer(word, letter, count):\n",
        "    parts = [\n",
        "        \"Found\",\n",
        "        f\"{count:04d}\",\n",
        "        \"occurrences of the letter\",\n",
        "        letter,\n",
        "        \"in\",\n",
        "        word\n",
        "    ]\n",
        "    return \" \".join(parts)\n",
        "\n",
        "def f(c):  # << START_OF_TRACE\n",
        "    word = \"strawberry\"\n",
        "    num = count_letters(word, c)\n",
        "    ans = format_answer(word, c, num)\n",
        "    return ans\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e959df87-60dc-4c32-8798-57914657dd3d",
      "metadata": {},
      "outputs": [],
      "source": [
        "idbg = InteractiveDebugger(code)\n",
        "idbg.display()"
      ]
    }
  ],
  "metadata": {
    "fileHeader": "",
    "fileUid": "4f384cd0-8315-48b4-a5d6-6d7047368aff",
    "isAdHoc": false,
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
